package com.dbms.storage;

import com.dbms.common.RID;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

/**
 * B+树实现类，用于高效存储和索引键值对（键为Integer，值为RID）。
 * 特性：叶子节点存储键与RID映射，内部节点存储键与子页面ID映射；
 * 依赖缓冲池管理器（BufferPoolManager）管理页面的加载、卸载和修改；
 * 支持插入、查找、删除（简化版）和树结构打印操作。
 */
public class BPlusTree {

    /** 缓冲池管理器，负责B+树页面的获取、释放、unpin（解除固定）和脏页标记 */
    private final BufferPoolManager bufferPoolManager;
    /** B+树的根页面ID，INVALID_PAGE_ID表示树为空 */
    private int rootPageId;
    /** 叶子节点的最大容量（最大可存储的键值对数量） */
    private final int leafMaxSize;
    /** 内部节点的最大容量（最大可存储的键-子页面ID对数量） */
    private final int internalMaxSize;

    /**
     * B+树构造方法，初始化树的核心依赖和参数
     * @param bufferPoolManager 缓冲池管理器，用于页面管理
     * @param rootPageId 初始根页面ID，空树时为INVALID_PAGE_ID
     * @param leafMaxSize 叶子节点最大容量
     * @param internalMaxSize 内部节点最大容量
     */
    public BPlusTree(BufferPoolManager bufferPoolManager, int rootPageId, int leafMaxSize, int internalMaxSize) {
        this.bufferPoolManager = bufferPoolManager;
        this.rootPageId = rootPageId;
        this.leafMaxSize = leafMaxSize;
        this.internalMaxSize = internalMaxSize;
    }

    /**
     * 【新增】获取 B+ 树的根页面 ID，用于外部访问树的根节点
     * @return 根页面 ID（INVALID_PAGE_ID 表示树为空）
     */
    public int getRootPageId() {
        return rootPageId;
    }

    /**
     * 判断B+树是否为空
     * @return true 表示树空（根页面ID为无效值），false 表示树非空
     */
    public boolean isEmpty() {
        return rootPageId == BPlusTreePage.INVALID_PAGE_ID;
    }

    /**
     * 根据键（key）查找对应的记录标识符（RID）
     * @param key 要查找的键
     * @return 找到的RID；若树空或键不存在，返回null
     * @throws IOException 缓冲池操作（如获取页面）可能抛出IO异常
     */
    public RID find(Integer key) throws IOException {
        // 树空直接返回null
        if (isEmpty()) return null;

        // 1. 找到包含该key的叶子节点页面
        Page leafPageObject = findLeafPage(key);
        // 将通用Page包装为B+树叶子节点页面
        BPlusTreeLeafPage leafPage = new BPlusTreeLeafPage(leafPageObject);

        // 2. 在叶子节点中查找key的索引（二分查找或线性查找，取决于leafPage.keyIndex实现）
        int index = leafPage.keyIndex(key);
        RID result = null;

        // 3. 验证索引有效性（是否在节点容量内）且对应键匹配，匹配则获取RID
        if (index < leafPage.getSize() && leafPage.getKey(index).equals(key)) {
            result = leafPage.getRid(index);
        }

        // 4. 解除叶子页面的固定（查询操作未修改页面，标记为非脏页）
        bufferPoolManager.unpinPage(leafPageObject.getPageId(), false);
        return result;
    }

    /**
     * 插入键值对（key-RID）到B+树中
     * 逻辑：找到对应叶子节点 → 插入键值对 → 若节点超容则分裂 → 分裂后更新父节点
     * @param key 要插入的键
     * @param rid 键对应的记录标识符
     * @throws IOException 缓冲池操作或节点分裂可能抛出IO异常
     */
    public void insert(Integer key, RID rid) throws IOException {
        // 情况1：树空，初始化新树
        if (isEmpty()) {
            startNewTree(key, rid);
            return;
        }

        // 情况2：树非空，找到插入目标叶子节点
        Page leafPageObject = findLeafPage(key);
        BPlusTreeLeafPage leafPage = new BPlusTreeLeafPage(leafPageObject);

        // 检查key是否已存在（避免重复插入），存在则直接unpin页面并返回
        if (leafPage.keyIndex(key) < leafPage.getSize() && leafPage.getKey(leafPage.keyIndex(key)).equals(key)) {
            bufferPoolManager.unpinPage(leafPage.getPageId(), false);
            return; // Key already exists
        }

        // 3. 在叶子节点的正确位置插入键值对（保持键有序）
        leafPage.insertAt(leafPage.keyIndex(key), key, rid);

        // 4. 若叶子节点未超容，直接unpin（标记为脏页，因内容已修改）并返回
        if (leafPage.getSize() < leafMaxSize) {
            bufferPoolManager.unpinPage(leafPage.getPageId(), true);
            return;
        }

        // 5. 叶子节点超容，执行分裂操作，生成新叶子节点
        BPlusTreeLeafPage newLeafPage = split(leafPage);
        // 6. 将新叶子节点的最小键（第一个键）插入到父节点中，建立父子关联
        insertIntoParent(leafPage, newLeafPage.getKey(0), newLeafPage);

        // 7. 解除原叶子节点和新叶子节点的固定（均已修改，标记为脏页）
        bufferPoolManager.unpinPage(leafPage.getPageId(), true);
        bufferPoolManager.unpinPage(newLeafPage.getPageId(), true);
    }

    /**
     * 初始化新的B+树（树空时调用）
     * 逻辑：创建根页面 → 初始化根为叶子节点 → 插入第一个键值对
     * @param key 第一个插入的键
     * @param rid 第一个键对应的RID
     * @throws IOException 缓冲池创建新页面可能抛出IO异常
     */
    private void startNewTree(Integer key, RID rid) throws IOException {
        // 1. 从缓冲池申请一个新页面作为根页面
        Page rootPageObject = bufferPoolManager.newPage();
        this.rootPageId = rootPageObject.getPageId(); // 更新根页面ID

        // 2. 将新页面包装为叶子节点，并初始化（设置最大容量、父页面ID等）
        BPlusTreeLeafPage rootPage = new BPlusTreeLeafPage(rootPageObject);
        rootPage.init(leafMaxSize); // 初始化叶子节点的元数据（如容量、空状态）
        rootPage.setParentPageId(BPlusTreePage.INVALID_PAGE_ID); // 根节点无父节点

        // 3. 插入第一个键值对（索引0，因节点为空）
        rootPage.insertAt(0, key, rid);

        // 4. 解除根页面固定（已修改，标记为脏页）
        bufferPoolManager.unpinPage(rootPageId, true);
    }

    /**
     * 分裂B+树节点（支持叶子节点和内部节点）
     * 逻辑：创建新节点 → 原节点将一半数据移动到新节点 → 维护节点间关联（如叶子的next/prev）
     * @param node 要分裂的节点（叶子或内部节点）
     * @return 分裂后生成的新节点
     * @throws IOException 缓冲池创建新页面可能抛出IO异常
     */
    private <T extends BPlusTreePage> T split(T node) throws IOException {
        // 1. 从缓冲池申请新页面，用于存储分裂后的部分数据
        Page newPageObject = bufferPoolManager.newPage();
        T newNode; // 分裂后生成的新节点

        // 情况1：分裂的是叶子节点
        if (node.getPageType() == BPlusTreePage.PageType.LEAF_PAGE) {
            BPlusTreeLeafPage leafNode = (BPlusTreeLeafPage) node;
            BPlusTreeLeafPage newLeafNode = new BPlusTreeLeafPage(newPageObject);

            // 初始化新叶子节点（设置最大容量）
            newLeafNode.init(leafMaxSize);
            // 原叶子节点将一半数据移动到新叶子节点（保持键有序）
            leafNode.moveHalfTo(newLeafNode);

            // 维护叶子节点的双向链表（更新next/prev指针，保证顺序访问）
            newLeafNode.setNextPageId(leafNode.getNextPageId());
            leafNode.setNextPageId(newLeafNode.getPageId());
            newLeafNode.setPrevPageId(leafNode.getPageId());

            newNode = (T) newLeafNode; // 泛型转换，适配返回类型
        }
        // 情况2：分裂的是内部节点
        else {
            BPlusTreeInternalPage internalNode = (BPlusTreeInternalPage) node;
            BPlusTreeInternalPage newInternalNode = new BPlusTreeInternalPage(newPageObject);

            // 初始化新内部节点（设置最大容量）
            newInternalNode.init(internalMaxSize);
            // 原内部节点将一半数据移动到新内部节点
            internalNode.moveHalfTo(newInternalNode);

            newNode = (T) newInternalNode; // 泛型转换，适配返回类型
        }

        // 设置新节点的父节点ID（与原节点相同，后续insertIntoParent会调整）
        newNode.setParentPageId(node.getParentPageId());
        return newNode;
    }

    /**
     * 将分裂后的新节点信息插入到父节点中（核心父节点更新逻辑）
     * 逻辑：无父节点（根分裂）→ 新建根内部节点；有父节点→插入父节点，超容则递归分裂
     * @param oldNode 分裂前的原节点（新节点的兄弟节点）
     * @param key 用于索引新节点的键（新节点的最小键）
     * @param newNode 分裂生成的新节点
     * @throws IOException 缓冲池操作或递归分裂可能抛出IO异常
     */
    private void insertIntoParent(BPlusTreePage oldNode, Integer key, BPlusTreePage newNode) throws IOException {
        // 情况1：原节点无父节点（说明原节点是根，分裂后需新建根内部节点）
        if (oldNode.getParentPageId() == BPlusTreePage.INVALID_PAGE_ID) {
            // 1. 申请新页面作为新根（内部节点）
            Page newRootPageObject = bufferPoolManager.newPage();
            BPlusTreeInternalPage newRoot = new BPlusTreeInternalPage(newRootPageObject);

            // 2. 初始化新根：设置最大容量，填充原节点、键、新节点的映射
            newRoot.init(internalMaxSize);
            newRoot.populateNewRoot(oldNode.getPageId(), key, newNode.getPageId());

            // 3. 更新原节点和新节点的父节点ID为新根的ID
            oldNode.setParentPageId(newRoot.getPageId());
            newNode.setParentPageId(newRoot.getPageId());

            // 4. 更新B+树的根页面ID为新根ID
            this.rootPageId = newRoot.getPageId();

            // 5. 解除新根页面固定（已修改，标记为脏页）
            bufferPoolManager.unpinPage(newRoot.getPageId(), true);
            return;
        }

        // 情况2：原节点有父节点，获取父节点页面
        Page parentPageObject = bufferPoolManager.fetchPage(oldNode.getParentPageId());
        BPlusTreeInternalPage parentNode = new BPlusTreeInternalPage(parentPageObject);

        // 子情况2.1：父节点未超容，直接插入新键和新节点的映射
        if (parentNode.getSize() < internalMaxSize) {
            // 在父节点中找到原节点的位置，在其后插入新键和新节点ID
            parentNode.insertNodeAfter(oldNode.getPageId(), key, newNode.getPageId());
            // 更新新节点的父节点ID（确保与父节点一致）
            newNode.setParentPageId(parentNode.getPageId());

            // 解除父节点固定（已修改，标记为脏页）
            bufferPoolManager.unpinPage(parentNode.getPageId(), true);
            return;
        }

        // 子情况2.2：父节点超容，需分裂父节点（递归处理）
        // 步骤1：创建临时列表，存储父节点的所有键和子页面ID（便于分裂后重新分配）
        List<Integer> tempKeys = new ArrayList<>();
        List<Integer> tempPointers = new ArrayList<>();
        // 注意：内部节点的键从索引1开始（索引0为第一个子页面ID，无对应键）
        for (int i = 0; i < parentNode.getSize(); i++) {
            if (i > 0) tempKeys.add(parentNode.getKey(i)); // 键从i=1开始添加
            tempPointers.add(parentNode.getPageId(i)); // 子页面ID从i=0开始添加
        }

        // 步骤2：找到新键和新子页面ID在临时列表中的插入位置（保持键有序）
        int insertIndex = 0;
        while (insertIndex < tempKeys.size() && tempKeys.get(insertIndex) < key) {
            insertIndex++;
        }
        tempKeys.add(insertIndex, key); // 插入新键
        tempPointers.add(insertIndex + 1, newNode.getPageId()); // 新子页面ID插入到键的下一位

        // 步骤3：分裂父节点，生成新父节点
        BPlusTreeInternalPage newParentNode = split(parentNode);

        // 步骤4：清空原父节点，重新填充前半部分数据（分裂后原父节点存储左半部分）
        parentNode.setSize(0); // 重置原父节点大小
        int middleIndex = tempKeys.size() / 2; // 中间键的索引（用于提升为父节点的键）
        Integer middleKey = tempKeys.get(middleIndex); // 分裂后需插入到祖父节点的键

        // 填充原父节点的第一个子页面ID
        parentNode.setPageId(0, tempPointers.get(0));
        parentNode.setSize(1); // 初始大小为1（仅第一个子页面ID）
        // 填充前半部分的键和子页面ID（从0到middleIndex-1）
        for (int i = 0; i < middleIndex; i++) {
            parentNode.setKey(i + 1, tempKeys.get(i)); // 键存储在i+1位置（对应i位置的子页面ID）
            parentNode.setPageId(i + 1, tempPointers.get(i + 1));
            parentNode.increaseSize(1); // 每填充一个键值对，大小+1
        }

        // 步骤5：填充新父节点，存储后半部分数据（从middleIndex+1开始）
        newParentNode.setPageId(0, tempPointers.get(middleIndex + 1)); // 新父节点的第一个子页面ID
        newParentNode.setSize(1); // 初始大小为1
        // 填充后半部分的键和子页面ID（从middleIndex+1到tempKeys.size()-1）
        for (int i = middleIndex + 1; i < tempKeys.size(); i++) {
            newParentNode.setKey(newParentNode.getSize(), tempKeys.get(i));
            newParentNode.setPageId(newParentNode.getSize(), tempPointers.get(i + 1));
            newParentNode.increaseSize(1);
        }

        // 【关键修复】更新新父节点所有子节点的父ID（确保子节点指向正确的父节点）
        updateChildrenParentIds(newParentNode);

        // 步骤6：递归调用，将中间键和新父节点插入到祖父节点中
        insertIntoParent(parentNode, middleKey, newParentNode);

        // 步骤7：解除原父节点和新父节点的固定（均已修改，标记为脏页）
        bufferPoolManager.unpinPage(parentNode.getPageId(), true);
        bufferPoolManager.unpinPage(newParentNode.getPageId(), true);
    }

    /**
     * 更新内部节点所有子节点的父页面ID（确保子节点与父节点关联正确）
     * 场景：父节点分裂后，新父节点的子节点需更新父ID为新父节点的ID
     * @param node 需更新子节点父ID的内部节点
     * @throws IOException 缓冲池获取子页面可能抛出IO异常
     */
    private void updateChildrenParentIds(BPlusTreeInternalPage node) throws IOException {
        // 遍历内部节点的所有子页面ID
        for (int i = 0; i < node.getSize(); i++) {
            int childPageId = node.getPageId(i);
            // 从缓冲池获取子页面
            Page childPageObject = bufferPoolManager.fetchPage(childPageId);
            // 将子页面转换为BPlusTreePage（叶子或内部节点）
            BPlusTreePage childPage = castToBPlusTreePage(childPageObject);
            // 更新子节点的父页面ID为当前内部节点的ID
            childPage.setParentPageId(node.getPageId());
            // 解除子页面固定（已修改父ID，标记为脏页）
            bufferPoolManager.unpinPage(childPageId, true);
        }
    }

    /**
     * 找到包含指定key的叶子节点页面（核心查找路径逻辑）
     * 逻辑：从根节点开始，递归（此处用循环）向下遍历内部节点，直到叶子节点
     * @param key 目标键
     * @return 包含该key的叶子节点页面（未unpin，需调用者后续处理）
     * @throws IOException 缓冲池获取页面可能抛出IO异常
     */
    private Page findLeafPage(Integer key) throws IOException {
        int currentPageId = rootPageId; // 从根节点开始
        // 从缓冲池获取当前页面
        Page currentPageObject = bufferPoolManager.fetchPage(currentPageId);
        // 转换为BPlusTreePage，判断节点类型
        BPlusTreePage currentPage = castToBPlusTreePage(currentPageObject);

        // 循环向下遍历：若当前节点不是叶子，继续查找子节点
        while (currentPage.getPageType() != BPlusTreePage.PageType.LEAF_PAGE) {
            // 转换为内部节点（因已判断非叶子）
            BPlusTreeInternalPage internalPage = (BPlusTreeInternalPage) currentPage;
            // 查找下一个子节点的页面ID（根据key在内部节点中的位置）
            int next_page_id = internalPage.lookUp(key);

            // 解除当前内部节点的固定（查询未修改，标记为非脏页）
            bufferPoolManager.unpinPage(currentPageId, false);

            // 更新当前节点信息，准备下一次循环
            currentPageId = next_page_id;
            currentPageObject = bufferPoolManager.fetchPage(currentPageId);
            currentPage = castToBPlusTreePage(currentPageObject);
        }

        // 循环结束，当前页面为叶子节点，返回（未unpin，需调用者处理）
        return currentPageObject;
    }

    /**
     * 将通用Page对象转换为BPlusTreePage（叶子节点或内部节点）
     * 依据：Page字节数据中PAGE_TYPE_OFFSET位置存储的页面类型标识（0=叶子，非0=内部）
     * @param page 待转换的通用Page对象
     * @return 转换后的BPlusTreePage（叶子或内部节点）
     */
    private BPlusTreePage castToBPlusTreePage(Page page) {
        // 包装Page的字节数据为ByteBuffer，便于读取特定位置的值
        ByteBuffer buffer = ByteBuffer.wrap(page.getData());
        // 读取PAGE_TYPE_OFFSET位置的页面类型标识
        int pageType = buffer.getInt(BPlusTreePage.PAGE_TYPE_OFFSET);

        // 根据页面类型创建对应的BPlusTreePage子类实例
        if (pageType == 0) { // 0 表示叶子节点（LEAF_PAGE）
            return new BPlusTreeLeafPage(page);
        } else { // 非0 表示内部节点（INTERNAL_PAGE）
            return new BPlusTreeInternalPage(page);
        }
    }

    /**
     * 【已实现】从 B+ 树中删除一个键（简化版）
     * 注意：简化点——仅删除叶子节点中的键值对，不处理节点容量不足时的合并/重分配逻辑
     * @param key 要删除的键
     * @throws IOException 缓冲池操作或叶子节点修改可能抛出IO异常
     */
    public void delete(Integer key) throws IOException {
        // 树空，直接返回（无键可删）
        if (isEmpty()) {
            return;
        }

        // 1. 找到包含该key的叶子节点页面
        Page leafPageObject = findLeafPage(key);
        BPlusTreeLeafPage leafPage = new BPlusTreeLeafPage(leafPageObject);

        // 2. 在叶子节点中查找key的索引
        int index = leafPage.keyIndex(key);

        // 3. 检查key是否存在：索引超出节点容量或键不匹配，直接unpin并返回
        if (index >= leafPage.getSize() || !leafPage.getKey(index).equals(key)) {
            bufferPoolManager.unpinPage(leafPage.getPageId(), false);
            return;
        }

        // 4. 从叶子节点中删除该键值对（保持剩余键有序）
        leafPage.removeAt(index);

        // 【简化点】完整B+树删除需补充：
        // - 检查叶子节点大小是否小于最小容量（minSize）
        // - 若小于，尝试与左/右兄弟节点合并，或从兄弟节点借数据
        // - 合并后若父节点键需删除，递归处理父节点的容量问题
        // 此处暂不实现上述逻辑

        // 5. 解除叶子节点固定（已修改，标记为脏页）
        bufferPoolManager.unpinPage(leafPage.getPageId(), true);
    }

    /**
     * 打印B+树的完整结构（用于调试和可视化）
     * 逻辑：调用私有递归打印方法，从根节点开始输出各节点信息
     * @throws IOException 缓冲池获取页面可能抛出IO异常
     */
    public void printTree() throws IOException {
        if (isEmpty()) {
            System.out.println("BPlusTree is empty.");
            return;
        }
        // 从根节点开始，层级0（无缩进）打印
        printTree(rootPageId, 0);
    }

    /**
     * 私有递归方法，打印指定页面及其子节点的信息
     * 格式：按层级缩进，输出节点类型、页面ID、大小、父ID和存储的键
     * @param pageId 当前要打印的页面ID
     * @param level 当前节点的层级（用于缩进，层级越高缩进越多）
     * @throws IOException 缓冲池获取页面可能抛出IO异常
     */
    private void printTree(int pageId, int level) throws IOException {
        // 无效页面ID，直接返回（递归终止条件之一）
        if (pageId == BPlusTreePage.INVALID_PAGE_ID) return;

        // 1. 获取当前页面并转换为BPlusTreePage
        Page pageObj = bufferPoolManager.fetchPage(pageId);
        BPlusTreePage page = castToBPlusTreePage(pageObj);

        // 2. 生成层级缩进（每级缩进2个空格，便于区分父子节点）
        String indent = "  ".repeat(level);

        // 3. 区分节点类型，打印不同信息
        if (page.getPageType() == BPlusTreePage.PageType.LEAF_PAGE) {
            BPlusTreeLeafPage leaf = (BPlusTreeLeafPage) page;
            // 叶子节点：输出类型、页面ID、大小、父ID、所有键
            System.out.print(indent + "Leaf Page " + pageId + " (size " + leaf.getSize() + ", parent " + leaf.getParentPageId() + "): ");
            for (int i = 0; i < leaf.getSize(); i++) {
                System.out.print(leaf.getKey(i) + " ");
            }
            System.out.println();
        } else {
            BPlusTreeInternalPage internal = (BPlusTreeInternalPage) page;
            // 内部节点：输出类型、页面ID、大小、父ID、所有键（跳过索引0的子页面ID对应键）
            System.out.print(indent + "Internal Page " + pageId + " (size " + internal.getSize() + ", parent " + internal.getParentPageId() + "): ");
            for (int i = 1; i < internal.getSize(); i++) {
                System.out.print(internal.getKey(i) + " ");
            }
            System.out.println();

            // 递归打印内部节点的所有子节点（层级+1）
            for (int i = 0; i < internal.getSize(); i++) {
                printTree(internal.getPageId(i), level + 1);
            }
        }

        // 4. 解除当前页面固定（打印未修改，标记为非脏页）
        bufferPoolManager.unpinPage(pageId, false);
    }

}